import numpy as np

def assign_cluster(X, centers):
    """将样本 X 分配到最近的簇中心"""
    # 计算所有样本到所有中心的距离
    distances = np.sqrt(((X[:, np.newaxis, :] - centers[np.newaxis, :, :]) ** 2).sum(axis=2))
    # 找到每个样本最近的簇中心索引
    labels = np.argmin(distances, axis=1)
    return labels


def kmeans(X, k, epsilon=1e-4, iteration=100):
    """手写实现 K-Means 算法"""
    # 随机初始化簇中心
    np.random.seed(42)
    indices = np.random.choice(X.shape[0], k, replace=False)
    centers = X[indices]

    for it in range(iteration):
        # 分配样本到最近的簇
        labels = assign_cluster(X, centers)

        # 更新每个簇的中心
        new_centers = np.zeros_like(centers)
        for j in range(k):
            points = X[labels == j]
            if len(points) > 0:
                new_centers[j] = np.mean(points, axis=0)
            else:
                new_centers[j] = centers[j]

        # 判断是否收敛
        shift = np.sqrt(np.sum((new_centers - centers) ** 2))
        if shift < epsilon:
            break  # 收敛，停止迭代

        centers = new_centers

    return centers, labels


if __name__ == "__main__":
    from sklearn.datasets import make_blobs
    import matplotlib.pyplot as plt

    # 生成测试数据集
    X, true_labels = make_blobs(n_samples=300, centers=3, cluster_std=0.7, random_state=42)

    # 执行 K-Means 聚类
    final_centroids, cluster_labels = kmeans(X, k=3, epsilon=1e-4, iteration=100)

    # 可视化聚类结果
    plt.figure(figsize=(10, 6))
    scatter = plt.scatter(X[:, 0], X[:, 1], c=cluster_labels, cmap='viridis', s=30, alpha=0.6)
    plt.scatter(final_centroids[:, 0], final_centroids[:, 1], 
                c='red', s=200, marker='X', linewidths=2, label='Cluster Centers')
    plt.colorbar(scatter, label='Cluster Labels')
    plt.title("K-Means Clustering Result")
    plt.xlabel("Feature 1")
    plt.ylabel("Feature 2")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()
